import struct
import mmap
import os
import fcntl
from PIL import Image
import numpy as np

class FramebufferReader:
    def __init__(self, fb_device='/dev/fb0'):
        self.fb_device = fb_device
        self.fd = None
        self.mm = None
        self.width = 0
        self.height = 0
        self.bpp = 0
        self.line_length = 0  # 行字节数（来自fix信息）
        self.smem_len = 0     # 总内存大小（来自fix信息）
        
    def open(self):
        try:
            # 1. 打开帧缓冲设备
            self.fd = os.open(self.fb_device, os.O_RDONLY)
            
            # 2. 获取固定信息（fix：line_length、内存大小等）
            # struct fb_fix_screeninfo 定义（内核标准结构）
            FBIOGET_FSCREENINFO = 0x4602
            # 格式：16字节id + 8个uint32 + 2个uint16 + 1个uint32 + 2个uint32 + 1个uint16 + 2个uint16
            fix_struct = struct.Struct('16sLIIIIHHHLLH2H')
            fix_buf = b'\x00' * fix_struct.size
            fix_data = fcntl.ioctl(self.fd, FBIOGET_FSCREENINFO, fix_buf)
            fix_info = fix_struct.unpack(fix_data)
            self.line_length = fix_info[8]  # line_length 在第9个字段（索引8）
            self.smem_len = fix_info[2]     # 总内存大小在第3个字段（索引2）
            
            # 3. 获取可变信息（var：分辨率、位深度等）
            FBIOGET_VSCREENINFO = 0x4600
            # struct fb_var_screeninfo 定义（简化关键字段，共40个uint32）
            var_struct = struct.Struct('I'*40)  # 严格匹配内核40个32位字段
            var_buf = b'\x00' * var_struct.size
            var_data = fcntl.ioctl(self.fd, FBIOGET_VSCREENINFO, var_buf)
            var_info = var_struct.unpack(var_data)
            self.width = var_info[0]    # xres（可见宽度）
            self.height = var_info[1]   # yres（可见高度）
            self.bpp = var_info[6]      # bits_per_pixel（位深度）
            
            # 4. 内存映射（使用实际内存大小smem_len）
            self.mm = mmap.mmap(self.fd, self.smem_len, prot=mmap.PROT_READ)
            
            print(f"✅ 成功打开帧缓冲区: {self.fb_device}")
            print(f"📊 分辨率: {self.width}x{self.height} | 位深度: {self.bpp}bpp")
            print(f"📏 行字节数: {self.line_length} | 总内存: {self.smem_len}字节")
            return True
            
        except PermissionError:
            print("❌ 权限不足！请用sudo运行")
        except FileNotFoundError:
            print(f"❌ 设备不存在: {self.fb_device}")
        except Exception as e:
            print(f"❌ 打开失败: {e}")
        finally:
            if not self.mm:
                self.close()
        return False
    
    def read_frame(self):
        if not self.mm:
            print("⚠️  帧缓冲区未打开")
            return None
            
        try:
            self.mm.seek(0)
            data = self.mm.read(self.smem_len)  # 读取全部内存数据
            
            if self.bpp == 32:
                # 尝试常见32位格式（按行字节数对齐）
                formats = ['BGRA', 'RGBA', 'ARGB', 'ABGR']
                for fmt in formats:
                    try:
                        img = Image.frombuffer(
                            'RGBA',
                            (self.width, self.height),
                            data,
                            'raw',
                            fmt,
                            self.line_length,  # 用实际行字节数对齐
                            1
                        )
                        return img.convert('RGB')
                    except:
                        continue
                print("❌ 32位格式解析失败")
                return None
                
            elif self.bpp == 24:
                try:
                    return Image.frombuffer(
                        'RGB',
                        (self.width, self.height),
                        data,
                        'raw',
                        'BGR',
                        self.line_length,
                        1
                    )
                except:
                    print("❌ 24位格式解析失败")
                    return None
                    
            elif self.bpp == 16:
                # 16位RGB565格式
                arr = np.frombuffer(data, dtype=np.uint16).reshape((self.height, self.width))
                r = ((arr >> 11) & 0x1F) << 3
                g = ((arr >> 5) & 0x3F) << 2
                b = (arr & 0x1F) << 3
                return Image.fromarray(np.dstack((r, g, b)).astype(np.uint8))
                
            else:
                print(f"⚠️  不支持的位深度: {self.bpp}bpp")
                return None
                
        except Exception as e:
            print(f"❌ 读取失败: {e}")
            return None
    
    def close(self):
        if self.mm:
            self.mm.close()
        if self.fd:
            os.close(self.fd)
        self.fd = self.mm = None
        print("🔌 帧缓冲区已关闭")

if __name__ == "__main__":
    fb_reader = FramebufferReader()
    if fb_reader.open():
        frame = fb_reader.read_frame()
        if frame:
            save_path = "fb_screenshot.png"
            frame.save(save_path)
            print(f"🖼️  截图已保存到 {save_path}")
            try:
                frame.show()
            except:
                print("⚠️  显示失败，可手动查看图片")
        fb_reader.close()